import numpy as np
from numpy.lib.function_base import gradient

import util

from linear_model import LinearModel


def main(train_path, eval_path, pred_path):
    """Problem 1(b): Logistic regression with Newton's Method.

    Args:
        train_path: Path to CSV file containing dataset for training.
        eval_path: Path to CSV file containing dataset for evaluation.
        pred_path: Path to save predictions.
    """
    x_train, y_train = util.load_dataset(train_path, add_intercept=True)

    # *** START CODE HERE ***
    #Train the model
    clf = LogisticRegression()
    clf.fit(x_train,y_train)

    #Load evaluation data
    x_eval, y_eval = util.load_dataset(eval_path, add_intercept=True)

    #Make Predictions
    predictions = clf.predict(x_eval)

    #Save predictions
    np.savetxt(pred_path, predictions)
    accuracy = np.mean(predictions == y_eval)
    print(f"P01b Accuracy: {accuracy:.4f}")

    # *** END CODE HERE ***


class LogisticRegression(LinearModel):
    """Logistic regression with Newton's Method as the solver.

    Example usage:
        > clf = LogisticRegression()
        > clf.fit(x_train, y_train)
        > clf.predict(x_eval)
    """

    def fit(self, x, y):
        """Run Newton's Method to minimize J(theta) for logistic regression.

        Args:
            x: Training example inputs. Shape (m, n).
            y: Training example labels. Shape (m,).
        """
        # *** START CODE HERE ***
        m, n =x.shape

        #initialize theta
        if self.theta is None:
            self.theta=np.zeros(n)

        #Newton's method iteration
        for iteration in range(self.max_iter):
            old_theta = self.theta.copy()

            #compute hypothesis h_theta(x)
            z = np.dot(x,self.theta)
            h_theta = 1 / (1 + np.exp(-z))

            #compute gradient
            gradient = -1/m * np.dot(x.T, (y-h_theta))

            #compute hessian & reshape h_theta for broadcasting
            h_theta_reshape = h_theta.reshape(-1, 1)
            #compute diagonal matrix D where D_ii = h_theta(x_i)(1-h_theta(x_i))
            D = np.diag(h_theta * (1- h_theta))
            hessian = 1/m * np.dot(np.dot(x.T, D), x)

            #update theta
            try:
                hessian_inv = np.linalg.inv(hessian)
                self.theta = self.theta - np.dot(hessian_inv, gradient)
            except np.linalg.LinAlgError:
                hessian_inv = np.linalg.pinv(hessian)
                self.theta = self.theta - np.dot(hessian_inv, gradient)

            #check for convergence
            if np.linalg.norm(self.theta - old_theta, 1)< self.eps:
                break
        # *** END CODE HERE ***

    def predict(self, x):
        """Make a prediction given new inputs x.

        Args:
            x: Inputs of shape (m, n).
        Returns:
            Outputs of shape (m,).
        """
        # *** START CODE HERE ***
        #compute hypothesis
        z = np.dot(x, self.theta)
        h_theta = 1/ (1 + np.exp(-z))

        return (h_theta >=0.5).astype(int)
        # *** END CODE HERE ***
